from astro.cosmology import Cosmology,PC
from astro.utilities import ismember, indexnear
from astro.io import readtxt
from astro.coord import ang_sep, match, s2dec
import astro.spec
import gas_gal_hiz.read
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.mlab import rec_append_fields

wlya = 1215.6701
wlyb = 1025.72
Ckms = 3.e5

GLOBAL_TRANS = 0.76

# adelberger data points and errors
A03 = readtxt('adel03-trans.dat', names='sep, fl, er')
A05 = readtxt('adel05-trans.dat', names='sep, fl, er')

cosmo = Cosmology(H0=100., Om=0.3, Ol=0.7)

# id numbers of qsos that were discovered with VIMOS (so are too faint
# to have aaomega spectra)
num_vimos = np.array([18, 19, 102, 103, 123, 124, 125, 126, 127,
                      143, 144, 145, 146])

# id numbers of qsos that are close enough to satisfy the angular
# spearation for LBG criteria, but don't actually have any LBGs with
# 15 Mpc of the forest, or only contribute a point or two at largest
# separations (>13 Mpc) and should be removed.
num_bad = np.array([1, 2, 3, 8, 11, 23, 55, 57, 58, 59, 61, 62, 63, 64, 68,
                    71, 72, 73, 76,  81, 83, 89, 90, 93, 94, 104, 111, 114,
                    115, 138])

# id numbers of qsos where we should use Williger's spectra rather
# than the AAOmega spectra.
num_will = np.array([1, 6, 7, 9, 10, 12, 14])

# id numbers and filenames of qsos observed with UVES/HIRES
# 131 to add in the PKS2126 QSO, 21_29_04.90m16_02_49.00

hires_filenames = {
    5   : '/home/nhmc/projects/gas_gal_hiz/qsos/Q0042m2627.dat',
    16  : '/home/nhmc/projects/gas_gal_hiz/qsos/Q0043m265cont.dat',
    60  : '/home/nhmc/projects/gas_gal_hiz/qsos/J0124p0044.asc',
    101 : '/home/nhmc/projects/gas_gal_hiz/qsos/HE0940m1050m.asc',
    110 : '/home/nhmc/projects/gas_gal_hiz/qsos/J1201p0116_hires.txt',
    131 : '/home/nhmc/projects/gas_gal_hiz/qsos/Q2129m1602.asc',
    134 : '/home/nhmc/projects/gas_gal_hiz/qsos/PKS2126m158.asc'}

num_hires = np.array(list(hires_filenames))

# masks that remove bad regions like poor sky subtraction, or DLAs.
MASKS = {
    16  : [(4625., 4660.)],                     # DLA
    60  : [(4920., 4993.)],                     # DLA
    84  : [(4730., 4775.)],                     # DLA
    110 : [(4392., 4546.)],                     # DLA
    131 : [(4630., 4665.)],                     # DLA
    141 : [(4725., 4770.), (5320., 5400.)]      # bad pixels
    }

# rough positions for the bright central QSOs in each field.

CRADECS = dict(HE0940  = s2dec('09 42 53.4','-11 04 25'),
               PKS2126 = s2dec('21 29 12.2','-15 38 41'),
               J0124   = s2dec('01 24 03.77','+00 44 32.7'),
               J1201   = s2dec('12 01 44.36',' +01 16 11.6'),
               Q0042   = s2dec('00 44 34.10',' -26 11 21.0'),
               OQ172   = s2dec('14 45 16.4652','+09 58 36.072'),
               Q2231   = s2dec('22 34 08.99','+00 00 01.7'),
               Q2359   = s2dec('00 01 40.6','+07 09 54'))

# plotting limits for transmissivity
YMIN, YMAX = -0.05, 1.2
XMIN, XMAX = -0.5,  15.0
A4LANDSCAPE = 11.7, 8.3
A4PORTRAIT = 8.3, 11.7

# maximum separation allowed for an LBG from a QSO sight-line
# (degrees).
MAXSEP = 15. / 60.

# bins to use when measuring mean flux. distances are in Mpc
maxdist = 20.0
binwidth = 0.5
BINEDGES = np.arange(0, maxdist + binwidth, binwidth)
CBINS = BINEDGES[:-1] + 0.5*binwidth
#binedges = list(binedges)

# hires_filenames_nok = {
#     5   : '/home/nhmc/projects/gas_gal_hiz/qsos/nok_qsos/fN01.dat',
#     16  : '/home/nhmc/projects/gas_gal_hiz/qsos/nok_qsos/fN02.dat',
#     60  : '/home/nhmc/projects/gas_gal_hiz/qsos/nok_qsos/fN03.dat',
#     101 : '/home/nhmc/projects/gas_gal_hiz/qsos/nok_qsos/fN04.dat',
#     110 : '/home/nhmc/projects/gas_gal_hiz/qsos/nok_qsos/fN05.dat',
#     134 : '/home/nhmc/projects/gas_gal_hiz/qsos/nok_qsos/fN06.dat'}

def readlbgs():
    # read LBG redshifts, ras and decs, find LBG comoving LOS
    # distances. (H0=100)
    lbgs = gas_gal_hiz.read.lbgs('all')

    # now only keep lbgs that are in rich's other list
    lbgs = readtxt(
        '/home/nhmc/projects/gas_gal_hiz/lbgs/LBGall_0.5_data_radecz.dat',
        names='num,ra,dec,z')
    #imatch = match(lbgs.ra,lbgs.dec, richlbgs.ra, richlbgs.dec, 2.0)
    const = 1. / (PC * 1.e6)
    losdist = [cosmo.Dc(z) * const  for z in lbgs.z]      # Mpc
    lbgs = rec_append_fields(lbgs, 'losdist', losdist)
    return lbgs

def read_randomlbgs():
    """ Read in the random LBG list.  Note there are 20x more random
    LBGs than the real catalogue."""
    const = 1. / (PC * 1.e6)
    rlbgs = readtxt(
        '/home/nhmc/projects/gas_gal_hiz/lbgs/LBGall_0.5_randoms_radecz.dat',
        names='ra,dec,z')
    num = range(len(rlbgs))
    rlbgs = rec_append_fields(rlbgs,'num', num)
    losdist = [cosmo.Dc(z) * const  for z in rlbgs.z]      # Mpc
    rlbgs = rec_append_fields(rlbgs, 'losdist', losdist)

    return rlbgs


def readqsos(zmin=2.0):
    """ Read in all the qsos and their appropriate spectra."""
    # read QSO ra, dec and emission redshift (note min(zqso) = 2.25)
    return gas_gal_hiz.read.qsos('all',zmin=zmin)


def filterqsos_for_lyalbg(qsos, lbgs):
    """ Keep only qsos that are close to LBGs have have either high or
    low res spectrum that is good.
    """
    # find subset of QSOs for cross correlation - all that are within
    # a few arcmin of an LBG.
    c0 = []
    for q in qsos:
        # separation is in degrees
        sep = ang_sep(q['ra'], q['dec'], lbgs.ra, lbgs.dec)
        minsep = sep.min()
        c0.append(True if minsep < MAXSEP else False)

    c0 = np.array(c0)
    # remove all that only have VIMOS spectra
    c1 = ~ismember(qsos.num, num_vimos)

    # and remove all that are 'bad' (not within 15 comoving Mpc of the
    # forest, or have low S/N, bad continuum or are a BALQSO)
    c2 = ~ismember(qsos.num, num_bad)

    qsos1 = qsos[c0 & c1 & c2]

    # take a look
    #plot(lbgs.ra,lbgs.dec,'g.')
    #plot(qsos1.ra,qsos1.dec,'o', mfc='None')
    #pl.axis('equal')
    # now find high resolution spectra
    is_hires_spec = ismember(qsos1.num, num_hires)
    qsos_hires = qsos1[is_hires_spec]
    # and aaomega spectra
    qsos_lores = qsos1[~is_hires_spec]
    print len(qsos_hires), len(qsos_lores)

    return qsos_hires, qsos_lores


def read_qso_spec(qsos):
    """ Read the spectra for each qso in an array read with
    readqsos()."""
    
    prefix = '/home/nhmc/projects/gas_gal_hiz/qsos/'
    aaomdirs = dict(Q0042   = prefix + 'aaomega/q0042/',
                    J0124   = prefix + 'aaomega/j0124/',
                    HE0940  = prefix + 'aaomega/he0940/',
                    PKS2126 = prefix + 'aaomega/pks2126/',
                    J1201   = prefix + 'aaomega/j1201/')    
    willigerdir = prefix + 'williger_qsos/'
    spec = []
    for qso in qsos:
        if qso['num'] in num_hires:
            filename = hires_filenames[qso['num']]
        elif qso['num'] in num_will:
            filename = willigerdir + qso['name'] + '.txt'
        else:
            filename = aaomdirs[qso['fieldqso']] + qso['name'] + '.txt'
        print 'Reading %s' % filename
        spec.append(astro.spec.read(filename))

    return spec


def calc_trans(binedges,
               lbgs_angdist, lbgs_losdist, lbgs_num,
               lyalosdist, nfl, qsonum, masked):
    """ Given a list of lbgs, LOS distances, normalised fluxes, masks
    and bins, returns mean flux per bin, mean flux per galaxy per bin,
    and the galaxy id numbers for each bin.
    """
    nbins = len(binedges) - 1
    
    # mean norm flux per galaxy per bin
    meanfl_gal_bins = [[] for i in range(nbins)] 
    # the id numbers of gals that contribute to each bin (for later checks)
    idnum = [[] for i in range(nbins)]
    # loop over all lbgs and bins.
    lenlbgs = len(lbgs_angdist)
    for j,(angdist, losdist) in enumerate(zip(lbgs_angdist,lbgs_losdist)):
        #print '%s of %s lbgs' % (j+1, lenlbgs)
        print '.',
        # separations between lbg position and all QSO pixel positions

        # note we have calculated the transverse comoving distance
        # (angdist) of an LBG from the QSO sightline at the redshift
        # of the LBG. This will introduce an error that increases with
        # distance from the

        sep = np.hypot(angdist, lyalosdist - losdist)
        for i in xrange(nbins):
            # find all pixel values that fall in this distance bin
            # that are not masked
            c0 = (binedges[i] <= sep) & (sep < binedges[i+1])
            nfl1 = nfl[c0 & ~masked]
            npix = len(nfl1)
            if npix > 0:
                #print 'contributes to bin %i' % i 
                meanfl_gal_bins[i].append(nfl1.mean())
                idnum[i].append( (lbgs_num[j], qsonum) )

    # no. of LBGs contributing each bin
    return meanfl_gal_bins, idnum

def trans_setup(qso, spec, lbgs, zscale=None, plot=False):
    """ Select appropiate QSO wav range and LBGs overlapping that
    range, Calculate comoving separation of LBGs from QSO sightline,
    and LOS comoving distance of the QSO pixels.

    Make a boolean mask to flag bad flux values.

    Scale the normalised flux if required.
    """
    # find minimum and maximum redshifts based on Lya forest region:
    dv = 3000.
    wmin = wlyb * (1.+ qso['z'])
    wmax = wlya * (1 + qso['z']) * (1 - dv/Ckms)

    dv = 5000.
    # go a little further beyond the z limits to look for LBGs.
    zmin = (wmin / wlya) * (1 - dv/Ckms) - 1
    zmax = (wmax / wlya) * (1 + dv/Ckms) - 1
    print 'Redshift limits %5.2f -> %5.2f' % (zmin,zmax)

    # lbg selection
    c0 = (zmin <= lbgs['z']) & (lbgs['z'] <= zmax)
    lbgs0 = lbgs[c0]
    sep = ang_sep(qso['ra'], qso['dec'], lbgs0.ra, lbgs0.dec)
    is_close =  sep < MAXSEP
    lbgs1 = lbgs0[is_close]
    # angular separation of lbgs from QSO sightline (radians)
    lbg_angsep = sep[is_close] * np.pi / 180.
    # comoving transverse separation 
    lbg_angdist = np.array([r*los for r,los in zip(lbg_angsep, lbgs1.losdist)])

    # only keep lbgs with separations that are smaller than our
    # largest separation bin limits.
    #c0 = lbg_angdist < max(binedges)
    #lbg_angdist = lbg_angdist[c0]
    #lbgs0 = lbgs0[c0]

    #if len(lbgs0) == 0:
    #    return (None,) * 5

    # select region of QSO spectrum we're interested in and normalise.
    spec0 = spec[(wmin <= spec.wa) & (spec.wa <= wmax)] 
    wa = spec0.wa
    nfl0 = spec0.fl / spec0.co

    ner = spec0.er / spec0.co
    zlya = wa / wlya - 1.
    const = 1 / (PC * 1.e6)
    lyalosdist = np.array([cosmo.Dc(z) * const for z in zlya]) 
    if zscale == False:
        nfl1 = nfl0  
    elif zscale == True:
        scaleval = GLOBAL_TRANS / (0.676 - 0.220*(zlya-3.))
        print 'Mean flux scaling value %.3f' % scaleval.mean()
        nfl1 = nfl0 * scaleval
    elif zscale == 'mean':
        good = (ner > 0) & ~np.isnan(nfl0)
        scaleval = GLOBAL_TRANS / nfl0[good].mean()
        print 'Flux scaling value %.3f' % scaleval
        nfl1 = nfl0 * scaleval
    else:
        raise ValueError('Unknown zscale value: %s' % zscale)
    
    # masking regions. Where masked=True, that pixel is not used.
    masked = np.ones(len(wa), dtype=bool)
    # unmask good points
    masked[(ner > 0) & ~np.isnan(nfl1)] = False
    if qso['num'] in MASKS:
        for wmin,wmax in MASKS[qso['num']]:
            print 'masking from %s to %s' % (wmin,wmax)
            condition = (wmin < wa) & ( wa < wmax)
            masked[condition] = True
        maskranges = MASKS[qso['num']]
    else:
        maskranges = None

    print '%i of %i  pixels masked' % (len(masked.nonzero()[0]), len(masked))

    if plot:
        plot_gal_lya(wa, nfl0, ner, lbgs1, lbg_angdist, qso['num'],
                     maskranges=maskranges)

    return lyalosdist, nfl1, masked, lbgs1, lbg_angdist


def run_trans_and_plot(qsos, spectra, lbgs, title=None, zscale=None,
                       per_spec_plots=False, singleplot=False,
                       info='hires', plotcheck=False):
    """ Calculate the mean flux near LBGs for a list of QSOs, and do
    some plots.
    """
    meanfl_per_gal_per_bin = [[] for i in range(len(CBINS))]
    idnum_per_bin = [[] for i in range(len(CBINS))]

    cols = 'brgycm'
    if zscale == 'mean': note = 'scalemean'
    elif zscale == True: note = 'scalefz'
    else: note = 'noscaling'

    if singleplot:
        fig = pl.figure(figsize=A4LANDSCAPE)
        ax = fig.add_subplot(111)
        ax.axhline(GLOBAL_TRANS, color='gray')

    for i,(qso,spec) in enumerate(zip(qsos, spectra)):
        print '\nCalculating trans for %(name)s' % qso 
        lyalosdist, nfl, masked, lbgs0, lbg_angdist = trans_setup(
            qso, spec, lbgs, zscale=zscale, plot=plotcheck)

        if lyalosdist is None:
            print 'There are no LBGs close enough to the forest, skipping'
            continue
        meanfl_per_gal, idnums = calc_trans(
            BINEDGES, lbg_angdist, lbgs0.losdist, lbgs0.num,
            lyalosdist, nfl, qso['num'], masked)
        
        if per_spec_plots:
            temp = '%(name)s_%(num)i' % qso
            filename = '%s_meanfl_%s_%s.pdf' % (info, temp, note)
            do_errplot(meanfl_per_gal, '%s %s' % (note, temp), filename)
        if singleplot:
            ngal = np.array([len(g) for g in idnums])
            sumfl = np.array([sum(f) for f in meanfl_per_gal])
            meanfl = sumfl / ngal
            c = cols[i % len(cols)]
            ax.plot(CBINS, meanfl, '-o' + c, label=qso['name'])

        # add the mean flux values and galaxy ids to a large list.
        for j in range(len(CBINS)):
            idnum_per_bin[j].extend(idnums[j])
            meanfl_per_gal_per_bin[j].extend(meanfl_per_gal[j])

    if singleplot:
        ax.set_title(title)
        ax.set_ylabel('Normalised Flux')
        ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
        ax.grid()
        ax.minorticks_on()
        ax.set_ylim(YMIN, YMAX)
        ax.set_xlim(XMIN, XMAX)
        if len(spectra) < 10:
            ax.legend(loc='lower right')
        fig.savefig('pdf/meanfl_%s_%s.pdf' % (info, note))        
        fig.savefig('meanfl_%s_%s.png' % (info, note))        
        
    #pl.show()

    return meanfl_per_gal_per_bin, idnum_per_bin


def do_boxplot(meanfl_per_gal, title, filename):
    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    ax.axhline(GLOBAL_TRANS, color='gray')
    ax.boxplot(meanfl_per_gal, positions=CBINS)
    ax.set_title(title)
    ax.set_ylabel('Normalised Flux')
    ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
    ax.grid()
    ax.minorticks_on()
    ax.set_ylim(YMIN, YMAX)
    ax.set_xlim(XMIN, XMAX)
    fig.savefig('pdf/' + filename)
    fig.savefig(filename.replace('.pdf','.png'))

def do_errplot(meanfl_per_gal, title, filename, show=False):
    """ Show the transmissivity as a function of distance between the
    lya forest and LBGs, with errors (standard error in the mean).
    """
    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    ax.axhline(GLOBAL_TRANS, color='gray')
    art = ax.errorbar(A03.sep, A03.fl, yerr=A03.er, fmt='o-k',mfc=None,
                      alpha=0.4)
    #for artist in art[:1]:
    #    artist.set_alpha(0.4)
    art = ax.errorbar(A05.sep, A05.fl, yerr=A05.er, fmt='ok',mfc=None,
                      alpha=0.4, linestyle='dotted')
    #for artist in art[:1]:
    #    artist.set_alpha(0.4)
    std =  np.array([np.std(f) for f in meanfl_per_gal])
    nvals = np.array([len(f) for f in meanfl_per_gal])
    stderr = std / np.sqrt(nvals)
    mean = [np.mean(f) for f in meanfl_per_gal]
    ax.errorbar(CBINS, mean, yerr=stderr, fmt='o-b')
    ax.set_title(title)
    ax.set_ylabel('Normalised Flux')
    ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
    ax.grid()
    ax.minorticks_on()
    ax.set_ylim(YMIN, YMAX)
    ax.set_xlim(XMIN, XMAX)
    fig.savefig('pdf/' + filename)
    fig.savefig(filename.replace('.pdf','.png'))
    if not show:
        pl.close()

def do_ngalplot(meanfl_per_gal, title, filename, show=False):
    """ Show the number of unique LBGs contributing to each bin.
    """
    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    ngal = [len(f) for f in meanfl_per_gal]
    pl.semilogy(CBINS, ngal, ls='steps-mid')
    ax.set_title(title)
    ax.set_ylabel('Number of galaxy-LOS pairs')
    ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
    ax.grid()
    ax.minorticks_on()
    ax.set_xlim(XMIN, XMAX)
    fig.savefig('pdf/' + filename)
    fig.savefig(filename.replace('.pdf','.png'))
    if not show:
        pl.close()

def plot_close_lbgs(allids, lbgs, qsos, spectra, dv=1500., lbgdv=490.):
    """ Plot the portions of QSO spectrum near LBGs less than 0.5, 1.0
    and 1.5 Mpc to the lya forest. This is wrapper around plot_spec_lbg().
    """
    ids = {}
    ids['<0.5'] = allids[0]
    ids['<1.0'] = set(allids[1]).difference(allids[0])
    ids['<1.5'] = set(allids[2]).difference(allids[0] + allids[1])

    for key in '<0.5 <1.0 <1.5'.split():
        if len(ids[key]) == 0:
            continue
        fig = pl.figure()
        fig.text(0.5, 0.92, key + ' $h^{-1}$ Mpc')
        for i,(idgal,idqso) in enumerate(ids[key]):
            if i == 0:
                ax = fig.add_subplot('%i1%i' % (len(ids[key]), i+1))
                ax0 = ax
            else:
                ax = fig.add_subplot('%i1%i' % (len(ids[key]), i+1))    
            plot_spec_lbg(idgal, idqso, lbgs, qsos, spectra,
                           ax=ax, dv=dv, lbgdv=lbgdv)

def plot_spec_lbg(idgal, idqso, lbgs, qsos, spectra, ax=None,
                   dv=1500., lbgdv=490.):
    """ Plot the portion of lya spectrum that is very close to an LBG.
    """
    import matplotlib.transforms as mtransforms
    if ax is None:
        fig = pl.figure()
        ax = fig.add_subplot(111)
    z, = lbgs[lbgs.num == idgal].z
    w0 = wlya * (1 + z)
    wmin = w0 * (1. - dv/Ckms)
    wmax = w0 * (1. + dv/Ckms)
    indqso = (qsos['num']==idqso).nonzero()[0][0]
    qso = qsos[indqso]
    spec = spectra[indqso]
    spec = spec[(wmin < spec.wa) & (spec.wa < wmax)]
    spec.multiply(1./spec.co)
    spec.plot(ax=ax)
    ax.axvline(w0,color='k', linestyle='dotted')
    trans = mtransforms.blended_transform_factory(
        ax.transData, ax.transAxes)
    x0,x1 = w0 * (1. - lbgdv / Ckms), w0 * (1 + lbgdv / Ckms)
    ax.fill([x0,x0,x1,x1],[0,1,1,0],fc='y', alpha=0.2,
            transform=trans)
    ax.text(0.95,0.75, '%(name)s\n%(num)i, z$_{em}=%(z).2f$' % qso,
            transform=ax.transAxes,ha='right')
    ax.text(0.05,0.85, 'LBG $z=%.2f$' % z,
            transform=ax.transAxes,ha='left')
    ax.set_xlim(wmin,wmax)
    ax.minorticks_on()

def plot_gal_lya(wa, nfl, ner, lbgs, lbg_angdist, num, maskranges=None):
    """
    Plot a lya spectrum and all the nearby LBGs. LBG distance from y=0
    is proportionalt to their distance from the QSO sight-line. y=1
    means 15 comoving h^-1 Mpc away.  A horizontal line marks the 2
    Mpc level.
    """
    if len(wa) == 0 or len(lbgs) == 0:
        print 'Either no wa values or no lbgs.'
        return
    ax = pl.gca()
    ax.cla()
    ax.axhline(0, color='0.7')
    ax.axhline(1/15.*2, color='0.7')
    ax.plot(wa, nfl, 'b', ls='steps-mid')
    ax.plot(wa, ner, color='orange')
    ax.scatter(wlya*(1 + lbgs.z), lbg_angdist / 15., marker='o',c='r')
    ax.text(0.9, 0.1, num, transform=ax.transAxes)
    if maskranges is not None:
        for w0,w1 in maskranges:
            y = -0.5, 1.5, 1.5,-0.5
            x = w0, w0, w1, w1
            ax.fill(x, y, alpha=0.4, fc='y')
    ax.set_ylim(-0.5,2)
    ax.set_xlim(wa[0], wa[-1])
    pl.show()
    raw_input('Enter')

def unique_absids(mat, ndim=1):
    """ Find the number of unique absorber ids in pair ID matrix.  If
    ndim == 2 assume a matrix of values, otherwise a 1d array."""
    len1 = len(mat)
    if ndim == 2:
        len2 = len(mat[0])
    absids = set([])
    for i in range(len1):
        if ndim == 2:
            for j in range(len2):
                temp = zip(*mat[i][j])
                if len(temp) > 1:
                    absids.update(temp[1])
        else:
            temp = zip(*mat[i])
            if len(temp) > 1:
                absids.update(temp[1])

    return np.array(sorted(absids))

def unique_lbgids(mat, ndim=1):
    """ Find the number of unique galaxy ids in pair ID matrix.  If
    ndim == 2 assume a matrix of values, otherwise a 1d array."""
    len1 = len(mat)
    if ndim == 2:
        len2 = len(mat[0])
    lbgids = set([])
    for i in range(len1):
        if ndim == 2:
            for j in range(len2):
                temp = zip(*mat[i][j])
                if len(temp) > 1:
                    lbgids.update(temp[0])
        else:
            temp = zip(*mat[i])
            if len(temp) > 1:
                lbgids.update(temp[0])
            
    return np.array(sorted(lbgids))

def plotfield(field, qsos, lbgs):
    """ Plot the qso and LBG ra/dec positions for a field, showing the
    numbers for each QSO.
    """
    cra, cdec = CRADECS[field.upper()]
    sep = ang_sep(lbgs.ra, lbgs.dec, cra, cdec)
    lbgs1 = lbgs[sep < 1]
    sep = ang_sep(qsos.ra, qsos.dec, cra, cdec)    
    qsos1 = qsos[sep < 1]

    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    ax.plot(lbgs1.ra, lbgs1.dec, '.b', alpha=0.7)
    ax.plot(qsos1.ra, qsos1.dec, 'rx', mew=1)
    ax.plot([cra],[cdec],'rx', mew=1, ms=10)
    for qso in qsos1:
        ax.text(qso['ra'], qso['dec'], qso['num'])
    ax.grid()
    ax.minorticks_on()
    ax.set_xlabel('RA (degrees)')
    ax.set_ylabel('Dec (degrees)')
    ax.set_title(field)
    ax.axis('equal')
    fig.savefig('pdf/radec_' + field.upper() + '.pdf')
    fig.savefig('radec_' + field.upper() + '.png')

if 0:
    lbgs = xcorr.readlbgs()
    qsos = xcorr.readqsos()
    qsos_hires, qsos_lores = xcorr.filterqsos_for_lyalbg(qsos, lbgs)
    spec_hires = xcorr.read_qso_spec(qsos_hires)
    spec_lores = xcorr.read_qso_spec(qsos_lores)

if 0:
    ###########################################################
    # Transmissivity calculation
    ###########################################################

#     meanfls, idgals = xcorr.run_trans_and_plot(
#         qsos_hires, spec_hires, lbgs,
#         title='No scaling', zscale=False,
#         singleplot=True, per_spec_plots=False)

    meanfls, idnums = xcorr.run_trans_and_plot(
        qsos_hires, spec_hires, lbgs,
        title='No scaling', zscale=False,
        singleplot=True, per_spec_plots=True, plotcheck=False)

    # Now plot the lot
    #xcorr.do_boxplot(meanfls,'6 HIRES QSOs', 'all_meanfl_hires_noscale.pdf')
    xcorr.do_errplot(meanfls,'6 HIRES QSOs no scaling',
                     'all_meanfl_hires_noscale_std.pdf')

    meanfls, idnums = xcorr.run_trans_and_plot(
        qsos_hires, spec_hires, lbgs,
        title='Scale flux by evolving expected mean', zscale=True,
        singleplot=True, per_spec_plots=False)

#     #xcorr.do_boxplot(meanfls,'6 HIRES QSOs',
#     #                 'all_meanfl_hires_scalefz.pdf')
    xcorr.do_errplot(meanfls,'6 HIRES QSOs scale by f(z)',
                     'all_meanfl_hires_scalefz_std.pdf')

    meanfls, idnums = xcorr.run_trans_and_plot(
        qsos_hires, spec_hires, lbgs,
        title='Scale flux by measured mean', zscale='mean',
        singleplot=True, per_spec_plots=False)

#     #xcorr.do_boxplot(meanfls, '6 HIRES QSOs',
#     #                 'all_meanfl_hires_scalemean.pdf')
    xcorr.do_errplot(meanfls, '6 HIRES QSOs scale by mean',
                     'all_meanfl_hires_scalemean_std.pdf')

    idnums_hires = idnums
    xcorr.do_ngalplot(idnums_hires, '6 HIRES', 'ngal_hires.pdf', show=True)

if 0:
#     meanfls, idnums = xcorr.run_trans_and_plot(
#         qsos_lores, spec_lores, lbgs,
#         title='No scaling', zscale=False, singleplot=False,
#         info='aaom', per_spec_plots=False, plotcheck=True)

#     #xcorr.do_boxplot(meanfls,'41 AA$\Omega$ QSOs',
#     #                 'all_meanfl_lores_noscale.pdf')
#     xcorr.do_errplot(meanfls, 'AA$\Omega$/Williger QSOs no scaling',
#                      'all_meanfl_lores_noscale_std.pdf')

    meanfls, idnums = xcorr.run_trans_and_plot(
        qsos_lores, spec_lores, lbgs,
        title='Scale flux by evolving expected mean', zscale=True,
        singleplot=True, info='aaom', per_spec_plots=False)

    #xcorr.do_boxplot(meanfls,'41 AA$\Omega$ QSOs',
    #                 'all_meanfl_lores_scalefz.pdf')
    xcorr.do_errplot(meanfls, 'AA$\Omega$/Williger QSOs scale by f(z)',
                     'all_meanfl_lores_scalefz_std.pdf')

#     meanfls, idnums = xcorr.run_trans_and_plot(
#         qsos_lores, spec_lores, lbgs,
#         title='Scale flux by measured mean', zscale='mean',
#         singleplot=True, info='aaom', per_spec_plots=False)

#     #xcorr.do_boxplot(meanfls,'41 AA$\Omega$ QSOs',
#     #                 'all_meanfl_lores_scalemean.pdf')
#     xcorr.do_errplot(meanfls, 'AA$\Omega$/Williger QSOs scale by mean',
#                      'all_meanfl_lores_scalemean_std.pdf')

    idnums_lores = idnums

    # plo of ngal for aaom and hires
    xcorr.do_ngalplot(idnums_lores, 'AAOm', 'ngal_lores.pdf', show=True)
    
if 0:
    # get ids of qsos used in xcorr
    idqsos_lo = unique_absids(idnums_lores)
    idqsos_hi = unique_absids(idnums_hires)
    
    # plot close lbgs
    qsos_all = np.concatenate([qsos_hires, qsos_lores])
    spec_all = np.concatenate([spec_hires, spec_lores])
    xcorr.plot_close_lbgs(idnums, lbgs, qsos_all, spec_all)

    # find lbgs within 1.5 h^-1 Mpc
    lbgsclose  = lbgs[ismember(lbgs.num, zip(*idnums[2])[0])]
    richlbgs = gas_gal_hiz.read.lbgs('all')
    imatch = match(lbgsclose.ra, lbgsclose.dec,
                   richlbgs.ra, richlbgs.dec, 3)
    #richlbgs[imatch] gives the same lbgs in rich's list.
